using BootstrapBlazorApp.Server.LLMs.api;
using BootstrapBlazorApp.Server.LLMs.chain;

namespace BootstrapBlazorApp.Server.LLMs.prompt;

public class ChatPromptTemplate : IBasePrompt
{
    private List<IBaseMessage> InputMessagesList { get; set; } = [];
    private List<IBaseMessage> Examples { get; init; } = [];

    public async Task<ChainContext> InvokeAsync(ChainContext input)
    {
        return await InvokeAsync(input.PromptInputValue,input);
    }

    public async IAsyncEnumerable<ChainContext> InvokeStreamAsync(ChainContext input)
    {
        yield return await InvokeAsync(input.PromptInputValue,input);
    }


    public static ChatPromptTemplate GetPromptTemplate(List<IBaseMessage> messages)
    {
        return new ChatPromptTemplate { InputMessagesList = messages };
    }

    public static ChatPromptTemplate FromMessages(List<IBaseMessage> messages,
        FewShotChatPromptTemplate fewShotChatPromptTemplate)
    {
        return new ChatPromptTemplate
        {
            InputMessagesList = messages, Examples = fewShotChatPromptTemplate.ToMessageList()
        };
    }


    public async Task<ChainContext> InvokeAsync(Dictionary<string, string> input)
    {
        
        foreach (var kvp in input)
        {
            foreach (var message in InputMessagesList.Where(message => message.Content.Contains($"{{{kvp.Key}}}")))
            {
                message.Content = message.Content.Replace($"{{{kvp.Key}}}", kvp.Value);
            }
        }
        Examples.AddRange(InputMessagesList);
        InputMessagesList = Examples;

        return await Task.FromResult(new ChainContext { MessagesList = InputMessagesList, PromptInputValue = input });
    }

    private async Task<ChainContext> InvokeAsync(Dictionary<string, string> input,ChainContext context)
    {
        
        foreach (var kvp in input)
        {
            foreach (var message in InputMessagesList.Where(message => message.Content.Contains($"{{{kvp.Key}}}")))
            {
                message.Content = message.Content.Replace($"{{{kvp.Key}}}", kvp.Value);
            }
        }
        Examples.AddRange(InputMessagesList);
        InputMessagesList = Examples;
        context.MessagesList.AddRange(InputMessagesList);
        context.PromptInputValue = input;
        return await Task.FromResult(context);
    }
}
